import glob
import random

import numpy as np
from torch.utils.data import Dataset
import torchvision.transforms as transforms
import os
from PIL import Image
import json
from pycocotools import mask as mask_utils
import torch
from torch.nn import functional as F
from tqdm import tqdm

def process_anns(anns, image_size, colormap):
    mask = np.zeros((image_size, image_size, 3))
    for i, ann in enumerate(anns):
        if ann['area'] < 5000:
            continue
        m = ann['segmentation']
        m = mask_utils.decode(m)
        X, Y = m.shape[1], m.shape[0]
        index = np.where(m == 1)
        x = int(np.mean(index[1]) // (X / 11))
        y = int(np.mean(index[0]) // (Y / 11))
        m = m.astype(bool)
        assert x * y < 124
        mask[m] = colormap[(x * y) % len(colormap)]
    return mask

def create_color_map():
    color_map = []
    for r in [0, 64, 128, 192, 255]:
        for g in [0, 64, 128, 192, 255]:
            for b in [0, 64, 128, 192, 255]:
                color_map.append([r, g, b])
    return np.array(color_map)[1:]


def find_classes(directory):
    """Finds the class folders in a dataset.

    See :class:`DatasetFolder` for details.
    """
    classes = sorted(entry.name for entry in os.scandir(directory) if entry.is_dir())
    if not classes:
        raise FileNotFoundError(f"Couldn't find any class folder in {directory}.")

    class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
    return classes, class_to_idx

class ImagenetC_gray_Dataset(Dataset):
    def __init__(self, root: str, split: str = "train", transform=None, image_size=256,
                 v_patch_nums=(1, 2, 3, 4, 5, 6, 8, 10, 13, 16), separator=False, val_cond='depth', **kwargs):
        self.transforms = transform
        self.split = split
        self.root = root
        
        # ImageNet 的标准目录结构
        self.split_dir = os.path.join(root, split)
        print('split_dir:', self.split_dir)
        
        # 获取所有类别文件夹
        classes, class_to_idx = find_classes(self.split_dir)
        self.classes = classes
        self.class_to_idx = class_to_idx
        
        # 获取所有图像路径
        self.image_paths = self._get_image_paths()
        print(f'Found {len(self.image_paths)} images in {self.split} split')
        
        self.cond_idx = {'gray': 5}  # 简化，只保留灰度 ，索引改为5
        self.v_patch_nums = v_patch_nums
        self.image_size = image_size
        self.separator = separator

    def _get_image_paths(self):
        """获取所有图像路径和对应的类别"""
        images = []
        for target_class in sorted(self.class_to_idx.keys()):
            class_index = self.class_to_idx[target_class]
            target_dir = os.path.join(self.split_dir, target_class)
            if not os.path.isdir(target_dir):
                continue
            for root, _, fnames in sorted(os.walk(target_dir, followlinks=True)):
                for fname in sorted(fnames):
                    if fname.lower().endswith(('.jpeg', '.jpg', '.png')):
                        path = os.path.join(root, fname)
                        images.append(path)
        return images

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, index):
        # 获取图像路径
        image_path = self.image_paths[index]
        
        # 获取类别
        class_name = os.path.basename(os.path.dirname(image_path))
        cls = self.class_to_idx[class_name]
        
        # 加载并转换图像
        try:
            image = Image.open(image_path).convert('RGB')
            image_gray = Image.open(image_path).convert('L').convert('RGB')
            
            # 调整图像大小
            if image.size != (self.image_size, self.image_size):
                image = image.resize((self.image_size, self.image_size))
                image_gray = image_gray.resize((self.image_size, self.image_size))
            
            if self.transforms:
                image, image_gray = self.transforms(image, image_gray)
            
            # 创建忽略掩码
            ignore_masks = torch.ones((1378,)) if self.separator else torch.ones((1360,))
            ignore_masks_ = torch.ones((1378,)) if self.separator else torch.ones((1360,))
            
            return {
                'image': image,
                'mask': image_gray,
                'cls': cls,
                'type': torch.tensor(self.cond_idx['gray']),
                'ignore_mask': ignore_masks,
                'ignore_mask_': ignore_masks_,
            }
            
        except Exception as e:
            print(f"Error loading image {image_path}: {str(e)}")
            # 返回数据集中的另一个样本
            return self.__getitem__((index + 1) % len(self))


if __name__ == '__main__':
    root= '../ImageNet2012/'
    cond_info_path = os.path.join(root, 'cond_info.json')
    if os.path.exists(cond_info_path):
        print('load ImageNetC from json')
        with open(cond_info_path, 'r') as file:
            cond_info = json.load(file)
            mask_paths = cond_info['mask']
            canny_paths = cond_info['canny']
            depth_paths = cond_info['depth']
            normal_paths = cond_info['normal']

    from tqdm import tqdm
    with tqdm(total=len(normal_paths)) as pbar:
        for path in normal_paths:
            try:
                img = Image.open(path)#.convert('RGB')
            except:
                print(path)
            # img_path = path.replace('normal', 'train').replace('.json', '.JPEG').replace('.jpeg', '.JPEG')
            #     with open(mask_path, 'r') as f:
            #         mask_info = json.load(f)
            #     mask = process_anns(mask_info, 512, colormap).astype(np.uint8)  # 512 is fixed during the labelling
            #     mask = Image.fromarray(mask)
            #     mask.save('mask.png')
            #     print(mask.path)
            # except:
            #     fail_list.append(mask_path)
            #     print(mask_path)

            pbar.update(1)
    # for path in fail_list:
    #     os.system(f'rm {path}')
